#!/usr/bin/env python3
"""
模型订阅客户端示例
演示如何使用API订阅模型更新并处理通知
"""

import asyncio
import json
import time
import requests
import websockets
from typing import Dict, Any
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelSubscriptionClient:
    """模型订阅客户端"""
    
    def __init__(self, base_url: str = "http://localhost:7880", client_id: str = "example-client"):
        self.base_url = base_url
        self.client_id = client_id
        self.client_name = f"示例客户端-{client_id}"
        self.session = requests.Session()
        self.subscriptions = []
        
    def get_models(self) -> Dict[str, Any]:
        """获取可用模型列表"""
        try:
            response = self.session.get(f"{self.base_url}/api/v1/models")
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"获取模型列表失败: {e}")
            return {"models": [], "total": 0}
    
    def get_servers(self) -> Dict[str, Any]:
        """获取服务器列表"""
        try:
            response = self.session.get(f"{self.base_url}/api/v1/servers")
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"获取服务器列表失败: {e}")
            return {"servers": [], "total": 0}
    
    def create_subscription(self, model_name: str, server_id: str, 
                          callback_url: str = None) -> Dict[str, Any]:
        """创建模型订阅"""
        data = {
            "client_id": self.client_id,
            "client_name": self.client_name,
            "model_name": model_name,
            "server_id": server_id,
            "auto_update": True,
            "notification_channels": ["webhook", "websocket"]
        }
        
        if callback_url:
            data["callback_url"] = callback_url
            
        try:
            response = self.session.post(
                f"{self.base_url}/api/v1/subscriptions",
                json=data
            )
            response.raise_for_status()
            result = response.json()
            
            if result.get("success"):
                logger.info(f"成功订阅模型: {model_name}")
                return result
            else:
                logger.error(f"订阅失败: {result.get('message')}")
                return result
                
        except Exception as e:
            logger.error(f"创建订阅失败: {e}")
            return {"success": False, "message": str(e)}
    
    def get_subscriptions(self) -> Dict[str, Any]:
        """获取当前客户端的订阅列表"""
        try:
            response = self.session.get(
                f"{self.base_url}/api/v1/subscriptions",
                params={"client_id": self.client_id}
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"获取订阅列表失败: {e}")
            return {"subscriptions": [], "total": 0}
    
    def delete_subscription(self, subscription_id: str) -> bool:
        """删除订阅"""
        try:
            response = self.session.delete(
                f"{self.base_url}/api/v1/subscriptions/{subscription_id}"
            )
            response.raise_for_status()
            result = response.json()
            
            if result.get("success"):
                logger.info(f"成功删除订阅: {subscription_id}")
                return True
            else:
                logger.error(f"删除订阅失败: {result.get('message')}")
                return False
                
        except Exception as e:
            logger.error(f"删除订阅失败: {e}")
            return False
    
    async def listen_websocket(self, duration: int = 60):
        """监听WebSocket消息"""
        ws_url = f"ws://localhost:7880/ws/{self.client_id}"
        
        try:
            logger.info(f"连接WebSocket: {ws_url}")
            async with websockets.connect(ws_url) as websocket:
                logger.info("WebSocket连接已建立")
                
                # 设置超时
                end_time = time.time() + duration
                
                while time.time() < end_time:
                    try:
                        # 等待消息，设置短超时以便检查时间
                        message = await asyncio.wait_for(
                            websocket.recv(), 
                            timeout=1.0
                        )
                        
                        data = json.loads(message)
                        await self.handle_websocket_message(data)
                        
                    except asyncio.TimeoutError:
                        # 超时是正常的，继续循环
                        continue
                    except websockets.exceptions.ConnectionClosed:
                        logger.warning("WebSocket连接已关闭")
                        break
                        
                logger.info("WebSocket监听结束")
                
        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}")
    
    async def handle_websocket_message(self, data: Dict[str, Any]):
        """处理WebSocket消息"""
        event_type = data.get("event_type", "unknown")
        
        logger.info(f"收到WebSocket消息: {event_type}")
        logger.info(f"消息内容: {json.dumps(data, indent=2, ensure_ascii=False)}")
        
        if event_type == "model_updated":
            await self.handle_model_update(data)
        elif event_type == "model_discovered":
            await self.handle_model_discovered(data)
        elif event_type == "subscription_created":
            await self.handle_subscription_created(data)
        else:
            logger.info(f"未处理的事件类型: {event_type}")
    
    async def handle_model_update(self, data: Dict[str, Any]):
        """处理模型更新事件"""
        model_name = data.get("model_name")
        model_id = data.get("model_id")
        
        logger.info(f"模型更新通知: {model_name} (ID: {model_id})")
        
        # 这里可以添加实际的模型更新处理逻辑
        # 例如：下载新模型、重新加载模型等
        
    async def handle_model_discovered(self, data: Dict[str, Any]):
        """处理模型发现事件"""
        model_name = data.get("model_name")
        server_id = data.get("server_id")
        
        logger.info(f"发现新模型: {model_name} (服务器: {server_id})")
        
    async def handle_subscription_created(self, data: Dict[str, Any]):
        """处理订阅创建事件"""
        logger.info("订阅创建通知")
    
    def run_demo(self):
        """运行演示"""
        logger.info("=== 模型订阅客户端演示 ===")
        
        # 1. 获取可用模型
        logger.info("1. 获取可用模型列表...")
        models_response = self.get_models()
        models = models_response.get("models", [])
        
        if not models:
            logger.warning("没有找到可用模型")
            return
            
        logger.info(f"找到 {len(models)} 个模型:")
        for model in models[:3]:  # 只显示前3个
            logger.info(f"  - {model['name']} v{model['version']} ({model['model_type']})")
        
        # 2. 获取服务器列表
        logger.info("\n2. 获取服务器列表...")
        servers_response = self.get_servers()
        servers = servers_response.get("servers", [])
        
        if not servers:
            logger.warning("没有找到可用服务器")
            return
            
        logger.info(f"找到 {len(servers)} 个服务器:")
        for server in servers:
            logger.info(f"  - {server['name']}: {server['url']} (状态: {server['status']})")
        
        # 3. 创建订阅
        logger.info("\n3. 创建模型订阅...")
        first_model = models[0]
        first_server = servers[0]
        
        subscription_result = self.create_subscription(
            model_name=first_model["name"],
            server_id=first_server["id"],
            callback_url="http://localhost:8080/webhook/model-update"
        )
        
        if not subscription_result.get("success"):
            logger.error("订阅创建失败")
            return
        
        # 4. 查看订阅列表
        logger.info("\n4. 查看当前订阅...")
        subscriptions_response = self.get_subscriptions()
        subscriptions = subscriptions_response.get("subscriptions", [])
        
        logger.info(f"当前有 {len(subscriptions)} 个订阅:")
        for sub in subscriptions:
            logger.info(f"  - {sub['model_name']} (状态: {sub['status']})")
        
        # 5. 监听WebSocket消息
        logger.info("\n5. 监听WebSocket消息 (30秒)...")
        logger.info("提示: 可以在另一个终端中触发事件来测试通知")
        logger.info("例如: curl -X POST http://localhost:7880/api/v1/discovery/trigger")
        
        try:
            asyncio.run(self.listen_websocket(duration=30))
        except KeyboardInterrupt:
            logger.info("用户中断WebSocket监听")
        
        # 6. 清理订阅（可选）
        logger.info("\n6. 清理演示订阅...")
        for sub in subscriptions:
            if sub["client_id"] == self.client_id:
                self.delete_subscription(sub["id"])
        
        logger.info("演示完成!")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="模型订阅客户端演示")
    parser.add_argument(
        "--base-url", 
        default="http://localhost:7880",
        help="API服务器地址"
    )
    parser.add_argument(
        "--client-id",
        default=f"demo-client-{int(time.time())}",
        help="客户端ID"
    )
    parser.add_argument(
        "--action",
        choices=["demo", "list-models", "list-servers", "list-subscriptions"],
        default="demo",
        help="执行的操作"
    )
    
    args = parser.parse_args()
    
    client = ModelSubscriptionClient(
        base_url=args.base_url,
        client_id=args.client_id
    )
    
    if args.action == "demo":
        client.run_demo()
    elif args.action == "list-models":
        models = client.get_models()
        print(json.dumps(models, indent=2, ensure_ascii=False))
    elif args.action == "list-servers":
        servers = client.get_servers()
        print(json.dumps(servers, indent=2, ensure_ascii=False))
    elif args.action == "list-subscriptions":
        subscriptions = client.get_subscriptions()
        print(json.dumps(subscriptions, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()
